import json
import os
import tempfile
from sys import exit
import lib.util

"""
文件格式
[文件格式标识(len:2)] +
[文件数量(len:2)] +
[ 
    [数据1长度(len:8)] 
    [数据1标识(len:2)] 
]  + 
[ 
    [数据2长度(len:8)] 
    [数据2标识(len:2)] 
] +
...
[ 
    [数据n长度(len:8)] 
    [数据n标识(len:2)] 
] +
[文件数据1(len:x)] +
[文件数据2(len:x)] +
....
[文件数据n(len:x)] +
[文件格式标识(2)]
"""


class Decode:
    fileName: str
    outDir: str
    password: str

    def __init__(self, fileName: str, outDir: str = "out", password=None):
        self.fileName = fileName
        if outDir is not None:
            self.outDir = outDir
        else:
            self.outDir = "out"
        self.password = password
        self.decode()

    # 文件解析
    def decode_file(self, temp_dir):
        with open(os.path.join(temp_dir, "data.temp"), 'r') as json_file:
            data_msg_temp = json.load(json_file)

        with open(os.path.join(temp_dir, "dp.temp"), "rb") as f:
            data_temp = f.read()

        data_flag_start = data_temp[:2]
        data_flag_end = data_temp[-2:]
        if data_flag_start != data_flag_end:
            raise Exception("The file is damaged and cannot be parsed")

        file_num = lib.util.bytes_to_int(data_temp[2:4])

        file_map = []
        for item in range(0, file_num):
            l, key, data = self.get_data(item, file_num, data_temp)
            file_map.append({
                "data": data,
                "len": l,
                "key": str(key),
                "file_name": data_msg_temp[str(key)]
            })
        return file_map

    # 获取文件信息
    def get_msg(self, index, data):
        # [标识(len:2)] + [文件数量(len:2)] + [ [数据长度(len:8)]  [数据标识(len:2)] ] +  [数据 len:n1] * n  +  [标识(2)]
        start = 4 + index * 10
        end = start + 10
        return lib.util.bytes_to_int(data[start: start + 8]), lib.util.bytes_to_int(data[start + 8: end])

    # 获取单个文件信息
    def get_data(self, index, msg_len, data):
        msg_data_len = 4 + msg_len * 10
        data_len = 0
        i_msg_data_len, i_msg_data_key = self.get_msg(index, data)
        for i in range(index):
            data_len += self.get_msg(i, data)[0]

        data_start = msg_data_len + data_len
        data_end = data_start + i_msg_data_len
        return i_msg_data_len, i_msg_data_key, data[data_start: data_end]

    def decode(self):
        # 创建临时文件夹 用来保存解包文件
        with tempfile.TemporaryDirectory() as tmp_dir:
            temp_dir = tmp_dir
            extract_flag = lib.util.extract_zip_file(self.fileName, temp_dir, self.password)
            if not extract_flag:
                print("Parsing failed please try again")
                exit(2)
            file_data = self.decode_file(temp_dir)
            if not os.path.exists(os.path.join(os.getcwd(), self.outDir)):
                os.makedirs(os.path.join(os.getcwd(), self.outDir))
            for item in file_data:
                with open(os.path.join(os.getcwd(), self.outDir, item["file_name"]), "wb") as f:
                    f.write(item["data"])
                print(f"write file [{item['file_name']}] success")
            print(f"out path is [{os.path.join(os.getcwd(), self.outDir)}]")
